import pandas as pd
import numpy as np
import talib
import yfinance as yf
import matplotlib.pyplot as plt

# MACD KDJ RSI B/S信号交易策略
class MACDKDJRSI:
    def __init__(self, period="1y", interval="1d"):
        self.period =period
        self.interval = interval

        self.get_data()
        self.calculate_all()
        self.marketing()

    def get_data(self):
        # 获取历史数据（示例用AAPL）
        # start = "2024-01-01"
        # end = "2025-02-26"
       self.df = yf.download(tickers="BTC-USD", period=self.period, interval=self.interval)
        # 1d season
        # 1h 3day
        # 30m day
        # 15m day

    # ----------------------------
    # 1. 计算MACD指标（参考‌:ml-citation{ref="4" data="citationList"}实现）
    # ----------------------------
    def calculate_macd(self):
        # 计算MACD指标
        # (
        #     self.df["MACD"],
        #     self.df["MACD_Signal_Line"],
        #     self.df["Histogram"],
        # ) = talib.MACD(
        #     self.df["Close"].values.flatten(), fastperiod=12, slowperiod=26, signalperiod=9
        # )

        self.df["EMA12"] = self.df["Close"].ewm(span=12, adjust=False).mean()
        self.df["EMA26"] = self.df["Close"].ewm(span=26, adjust=False).mean()
        self.df["DIF"] = self.df["EMA12"] - self.df["EMA26"]
        self.df["DEA"] = self.df["DIF"].ewm(span=9, adjust=False).mean()
        self.df["MACD"] = (self.df["DIF"] -self.df["DEA"]) * 2
        # return self.df


    # ----------------------------
    # 2. 计算KDJ指标（参考‌:ml-citation{ref="4,6" data="citationList"}实现）
    # ----------------------------
    def calculate_kdj(self, n=9, m1=3, m2=3):
        low_min =self.df["Low"].rolling(n, min_periods=1).min()
        high_max =self.df["High"].rolling(n, min_periods=1).max()
        self.df["RSV"] = (self.df["Close"] - low_min) / (high_max - low_min) * 100
        self.df["K"] =self.df["RSV"].ewm(alpha=1 / m1, adjust=False).mean()
        self.df["D"] =self.df["K"].ewm(alpha=1 / m2, adjust=False).mean()
        self.df["J"] = 3 *self.df["K"] - 2 *self.df["D"]


    # 3.计算RSI
    def calculate_rsi(self, window=14):
        # 提取收盘价序列
        close = self.df["Close"]

        # 计算价格变化差值
        delta = close.diff(1)

        # 分离涨跌幅
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)

        # 计算平均涨跌幅（使用指数移动平均）
        avg_gain = gain.ewm(alpha=1 / window, min_periods=window).mean()
        avg_loss = loss.ewm(alpha=1 / window, min_periods=window).mean()

        # 计算RS和RSI
        rs = avg_gain / avg_loss
        self.df["RSI"] = 100 - (100 / (1 + rs))


    # 执行计算
    def calculate_all(self):
        self.calculate_macd()
        self.calculate_kdj()
        self.calculate_rsi()
        self.df = self.df.fillna(0)


    def marketing(self):
        # ----------------------------
        # 3. 生成交易信号
        # ----------------------------
        # MACD金叉/死叉判断
       self.df["MACD_Cross"] = np.where(self.df["DIF"] >self.df["DEA"], 1, -1)
       self.df["MACD_Signal"] =self.df["MACD_Cross"].diff()

        # KDJ超买超卖判断（J>80超买，J<20超卖）
       self.df["KDJ_Over"] = np.where(self.df["J"] >self.df["D"], 1, -1)
        # data["KDJ_Over"] = (data["J"] > data["D"]) & data["J"]<50

        # 综合信号（MACD金叉+KDJ超卖时买入且RSI小于70时买入）
       self.df["Buy_Signal"] = (
            ((self.df["MACD_Signal"] > 0) & (self.df["KDJ_Over"] == 1))
            & (self.df["RSI"] < 80)
        )
       self.df["Sell_Signal"] = ((self.df["MACD_Signal"] < 0) & (self.df["KDJ_Over"] == -1)) & (
           self.df["RSI"] > 30
        ) 

        # print(data.head(10))

        # filter_data = self.df[(self.df["Buy_Signal"]) | (self.df["Sell_Signal"])]
        #
        # print(filter_data.head(100))

        # data.head(500)


    # ----------------------------
    # 4. 回测展示（简化版）
    # ----------------------------

    # 数据可视化
    def backtest(self):
        plt.figure(figsize=(16, 12))
        # 价格与MACD
        plt.subplot(4, 1, 1)
        plt.plot(self.df["Close"], label="Price")
        plt.plot(
           self.df[self.df["Buy_Signal"]].index,
           self.df["Close"][self.df["Buy_Signal"]],
            "^",
            markersize=10,
            color="r",
        )
        plt.plot(
           self.df[self.df["Sell_Signal"]].index,
           self.df["Close"][self.df["Sell_Signal"]],
            "v",
            markersize=10,
            color="g",
        )

        # MACD指标
        plt.subplot(4, 1, 2)
        plt.plot(self.df["DIF"], label="DIF")
        plt.plot(self.df["DEA"], label="DEA")
        plt.bar(
            self.df.index,
            self.df["MACD"],
            label="MACD Histogram",
            color=[("red" if x > 0 else "green") for x in self.df["MACD"]],
            alpha=0.5,
        )
        plt.legend(loc="lower left", frameon=False)
        # plt.bar(
        #     data.index,
        #     data["Histogram"],
        #     color=[("red" if x > 0 else "green") for x in data["Histogram"]],
        #     alpha=0.5,
        # )

        # KDJ指标
        plt.subplot(4, 1, 3)
        plt.plot(self.df["K"], label="K")
        plt.plot(self.df["D"], label="D")
        plt.plot(self.df["J"], label="J")
        plt.axhline(80, linestyle="--", color="gray")
        plt.axhline(20, linestyle="--", color="gray")
        plt.legend(loc="lower left", frameon=False)

        # RSI绘制
        plt.subplot(4, 1, 4)
        plt.plot(self.df["RSI"], color="navy")
        plt.fill_between(self.df.index, 30, 70, color="lightgrey", alpha=0.3)
        plt.axhline(70, color="red", linewidth=0.5)
        plt.axhline(30, color="green", linewidth=0.5)

        plt.show()


# Valid periods: 1d,5d,1mo,3mo,6mo,1y,2y,5y,10y,ytd,max
# Valid intervals: 1m,2m,5m,15m,30m,60m,90m,1h,1d,5d,1wk,1mo,3mo
if __name__ == '__main__':
    # MACDKDJRSI(period="3mo", interval="1d").backtest()
    MACDKDJRSI(period="1wk", interval="1h").backtest()
    # MACDKDJRSI(period="1d",interval="15m").backtest()

    import argparse

    # 创建解析器对象
    parser = argparse.ArgumentParser(description="-p periods; -i intervals")

    # 添加参数
    parser.add_argument(
        "-p", type=str, choices=["1d","5d","1mo","3mo","6mo","1y","2y","5y","10y","ytd","max"], help="periods"
    )
    parser.add_argument(
        "-i", type=str, choices=["1m","2m","5m","15m","30m","60m","90m","1h","1d","5d","1wk","1mo","3mo"], help="intervals"
    )

    # 解析参数
    args = parser.parse_args()

    period = "1y"
    interval = "1d"
    if args.p:
        period = args.p
    if args.i:
        interval = args.i
    
    # MACDKDJRSI(period=period, interval=interval).backtest()
